#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
二维码编码器模块
将文件数据编码为二维码序列
"""

import qrcode
import os
import struct
import zlib
from typing import List, Tuple


class QREncoder:
    # 压缩标志位
    COMPRESSION_FLAG = 0b00000001
    
    # 不需要压缩的文件扩展名列表
    NO_COMPRESSION_EXTENSIONS = {
        # 已压缩格式
        '.zip', '.rar', '.7z', '.gz', '.bz2', '.xz', '.tar', '.tgz',
        # 图片格式
        '.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp', '.ico', '.svg',
        # 音视频格式
        '.mp3', '.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.wav', '.flac', '.aac',
        # 其他二进制格式
        '.iso'
    }
    
    def __init__(self):
        # 二维码最大数据容量（字节）
        # 使用最低纠错级别L，版本40（177x177）
        # 二进制模式下理论最大容量为2953字节
        # 为内存和库内部开销预留空间，使用更保守的2900字节
        self.max_data_per_qr = 2900  # 保守估计每个二维码可承载的数据量
        self.enable_compression = True  # 默认启用压缩
    
    def _should_compress_file(self, filename: str) -> bool:
        """
        判断是否应该对文件进行压缩
        
        Args:
            filename: 文件名
            
        Returns:
            True: 应该压缩，False: 不应该压缩
        """
        if not self.enable_compression:
            return False
            
        # 获取文件扩展名
        _, ext = os.path.splitext(filename.lower())
        
        # 如果扩展名在不压缩列表中，则不压缩
        if ext in self.NO_COMPRESSION_EXTENSIONS:
            return False
            
        return True

    def estimate_qr_count(self, file_size: int, filename: str, compressed_size: int = None) -> int:
        """
        估算需要的二维码数量
        
        Args:
            file_size: 文件大小（字节）
            filename: 文件名
            compressed_size: 压缩后的大小（如果为None则使用原始大小）
            
        Returns:
            需要的二维码数量
        """
        # 计算文件名相关开销
        filename_bytes = filename.encode('utf-8')
        filename_length = len(filename_bytes)
        header_size = 1 + 1 + filename_length + 2 + 2 + 4  # 标志位1字节 + 文件名长度1字节 + 文件名 + 序号2字节 + 总数2字节 + 原始大小4字节
        payload_size = self.max_data_per_qr - header_size
        # 确保payload_size为正数
        if payload_size <= 0:
            raise ValueError(f"文件名过长，导致无法存储任何有效数据。文件名长度{len(filename_bytes)}字节，最大支持{self.max_data_per_qr-9}字节")

        # 根据文件扩展名决定是否使用压缩
        should_compress = self._should_compress_file(filename)
        
        # 如果文件不应该压缩，直接使用原始大小
        if not should_compress:
            data_size = file_size
        elif compressed_size is not None:
            # 使用提供的压缩大小
            data_size = compressed_size
        else:
            # 使用原始大小（压缩效果未知时）
            data_size = file_size
            
        qr_count = (data_size + payload_size - 1) // payload_size
        
        return max(qr_count, 1)
        
    def encode_file_to_qr(self, file_path: str, output_dir: str, progress_callback=None, logger=None) -> List[str]:
        """
        将文件编码为二维码序列（支持压缩）
        
        Args:
            file_path: 输入文件路径
            output_dir: 二维码输出目录
            progress_callback: 进度回调函数
            logger: 日志记录器
            
        Returns:
            生成的二维码文件路径列表
        """
        try:
            # 检查输出目录是否存在
            if not os.path.exists(output_dir):
                try:
                    os.makedirs(output_dir)
                except OSError as e:
                    raise RuntimeError(f"无法创建输出目录 '{output_dir}': {str(e)}")
            elif not os.path.isdir(output_dir):
                raise RuntimeError(f"指定的输出路径 '{output_dir}' 不是目录")
                
            # 读取文件
            if not os.path.exists(file_path):
                raise RuntimeError(f"输入文件不存在: '{file_path}'")
            if not os.path.isfile(file_path):
                raise RuntimeError(f"指定路径不是文件: '{file_path}'")
                
            with open(file_path, 'rb') as f:
                file_data = f.read()
                
            filename = os.path.basename(file_path)
            file_size = len(file_data)
            
            # 检查空文件
            if file_size == 0:
                raise ValueError("无法处理空文件")
            
            # 判断是否应该压缩文件
            should_compress = self._should_compress_file(filename)
            
            if should_compress:
                # 尝试压缩文件数据
                try:
                    compressed_data = zlib.compress(file_data)
                    compressed_size = len(compressed_data)
                    use_compression = compressed_size < file_size  # 只在压缩后更小时使用
                    
                    if use_compression:
                        data_to_encode = compressed_data
                        actual_data_size = compressed_size
                        compression_ratio = (1 - compressed_size / file_size) * 100
                        if logger:
                            logger.success(f"[压缩] 启用压缩: {file_size} → {compressed_size} 字节 (压缩率: {compression_ratio:.1f}%)")
                    else:
                        data_to_encode = file_data
                        actual_data_size = file_size
                        use_compression = False
                        if logger:
                            logger.info("[压缩] 压缩无效，使用原始数据")
                        
                except Exception as e:
                    # 如果压缩失败，使用原始数据
                    if logger:
                        logger.warning(f"[压缩] 压缩失败: {e}，使用原始数据")
                    data_to_encode = file_data
                    actual_data_size = file_size
                    use_compression = False
            else:
                # 不压缩的文件类型
                data_to_encode = file_data
                actual_data_size = file_size
                use_compression = False
                if logger:
                    logger.info(f"[压缩] 跳过压缩: {filename} 为已压缩格式，直接使用原始数据")
                
            # 计算二维码数量（使用压缩后的大小）
            qr_count = self.estimate_qr_count(file_size, filename, actual_data_size)
            
            # 计算每个二维码的数据
            filename_bytes = filename.encode('utf-8')
            filename_length = len(filename_bytes)
            header_size = 1 + 1 + filename_length + 2 + 2 + 4  # 标志位1字节 + 文件名长度1字节 + 文件名 + 序号2字节 + 总数2字节 + 原始大小4字节
            payload_size = self.max_data_per_qr - header_size
            
            # 验证单个二维码数据量是否超出限制
            max_qr_capacity = self.max_data_per_qr
            chunk_data_size = min(payload_size, actual_data_size)
            total_chunk_size = header_size + chunk_data_size
            if total_chunk_size > max_qr_capacity:
                raise ValueError(f"单个二维码数据量({total_chunk_size}字节)超出最大容量限制({max_qr_capacity}字节)，请减小文件名长度或分割文件")
            
            qr_files = []
            
            # 生成每个二维码
            for i in range(qr_count):
                # 计算当前二维码的数据范围
                start_pos = i * payload_size
                end_pos = min((i + 1) * payload_size, actual_data_size)
                chunk_data = data_to_encode[start_pos:end_pos]

                # 构建二维码数据
                qr_data = self._build_qr_data(
                    filename_bytes, 
                    i, 
                    qr_count, 
                    chunk_data,
                    file_size,
                    use_compression
                )

                # 生成二维码
                try:
                    qr = qrcode.QRCode(
                        version=40,
                        error_correction=qrcode.constants.ERROR_CORRECT_L,
                        box_size=12,
                        border=8
                    )
                    qr.add_data(qr_data, optimize=0)
                    qr.make(fit=False)

                    # 保存二维码图片
                    img = qr.make_image(fill_color="black", back_color="white")
                    qr_file = os.path.join(output_dir, f"qr_{i:04d}.png")
                    img.save(qr_file)
                    qr_files.append(qr_file)

                    # 计算当前二维码的数据范围
                    start_pos = i * payload_size
                    end_pos = min((i + 1) * payload_size, actual_data_size)
                    chunk_size = end_pos - start_pos
                    
                    # 实时记录文件片段信息
                    if logger:
                        logger.info(f"文件片段 [{start_pos}:{end_pos}] ({chunk_size}字节) -> qr_{i:04d}.png")
                    
                    # 更新进度
                    if progress_callback:
                        progress = ((i + 1) / qr_count) * 100
                        progress_callback(progress)
                        
                except Exception as e:
                    raise RuntimeError(f"生成第{i+1}个二维码时出错: {str(e)}")
                
            return qr_files
            
        except MemoryError:
            raise RuntimeError("内存不足，无法处理大文件。请尝试减小文件大小或增加系统内存")
        except OSError as e:
            raise RuntimeError(f"文件操作错误: {str(e)}")
        except Exception as e:
            raise RuntimeError(f"二维码生成失败: {str(e)}")
        
    def _build_qr_data(self, filename_bytes: bytes, index: int, total: int, chunk_data: bytes, original_size: int, use_compression: bool) -> bytes:
        """
        构建二维码数据格式（支持压缩）
        
        格式：[标志位1字节][文件名长度1字节][文件名][序号2字节][总数2字节][原始大小4字节][文件内容]
        
        Args:
            filename_bytes: 文件名（UTF-8编码）
            index: 当前二维码序号（从0开始）
            total: 二维码总数
            chunk_data: 文件内容块（可能为压缩数据）
            original_size: 原始文件大小
            use_compression: 是否使用压缩
            
        Returns:
            完整的二维码数据
        """
        # 验证序号和总数的有效性
        if index < 0 or index > 65535:
            raise ValueError(f"二维码序号无效: {index}，必须在0-65535范围内")
        if total < 1 or total > 65535:
            raise ValueError(f"二维码总数无效: {total}，必须在1-65535范围内")
        if index >= total:
            raise ValueError(f"二维码序号{index}超出总数{total}")
            
        data = bytearray()
        
        # 标志位（1字节）
        flags = 0
        if use_compression:
            flags |= self.COMPRESSION_FLAG
        data.extend(struct.pack('B', flags))
        
        # 文件名长度（1字节）
        if len(filename_bytes) > 255:
            raise ValueError(f"文件名过长: {len(filename_bytes)}字节，最大支持255字节")
        data.extend(struct.pack('B', len(filename_bytes)))
        
        # 文件名
        data.extend(filename_bytes)
        
        # 序号（2字节，小端序，无符号短整型）
        data.extend(struct.pack('<H', index))
        
        # 总数（2字节，小端序，无符号短整型）
        data.extend(struct.pack('<H', total))
        
        # 原始文件大小（4字节，小端序，无符号整型）
        if original_size > 0xFFFFFFFF:
            raise ValueError(f"文件过大: {original_size}字节，最大支持4GB")
        data.extend(struct.pack('<I', original_size))
        
        # 文件内容
        data.extend(chunk_data)
        
        return bytes(data)